// MemSpyyDlg.cpp : implementation file
//

#include "stdafx.h"
#include "MemSpyy.h"
#include "MemSpyyDlg.h"
#include ".\MemSpyydlg.h"
#include "psapi.h"

#include "..\LibSpyDll\LibSpyDll.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// CAboutDlg dialog used for App About

class CAboutDlg : public CDialog
{
public:
	CAboutDlg();

// Dialog Data
	enum { IDD = IDD_ABOUTBOX };

	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV support

// Implementation
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialog(CAboutDlg::IDD)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialog)
END_MESSAGE_MAP()


// CMemSpyyDlg dialog



CMemSpyyDlg::CMemSpyyDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CMemSpyyDlg::IDD, pParent)
	, m_PID(0)
	, m_InjecteeProcess(NULL)
	, m_hLibModule(0)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CMemSpyyDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_LARGEST, m_LargestCtl);
	DDX_Text(pDX, IDC_PID, m_PID);
	DDX_Control(pDX, IDC_TOTALFREE, m_TotalFreeCtl);
	DDX_Control(pDX, IDC_OGLCUSTOM, m_OpenGLControl);
	DDX_Control(pDX, IDC_TOTALDLL, m_TotalDllCtl);
	DDX_Control(pDX, IDC_PID, m_PIDCtl);
	DDX_Control(pDX, IDC_PIDNAME, m_PIDNameCtl);
	DDX_Control(pDX, IDC_XCOORD, m_XCoordCtl);
	DDX_Control(pDX, IDC_YCOORD, m_YCoordCtl);
	DDX_Control(pDX, IDC_STATUS, m_StatusCtl);
	DDX_Control(pDX, IDC_PROCESSHEAP, m_ProcessHeapCtl);
	DDX_Control(pDX, IDC_TOTALVIRTUALSPACE, m_TotalVirtSpaceCtl);
}

BEGIN_MESSAGE_MAP(CMemSpyyDlg, CDialog)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	//}}AFX_MSG_MAP
	ON_WM_TIMER()
	ON_EN_UPDATE(IDC_PID, OnEnUpdatePid)
END_MESSAGE_MAP()


// CMemSpyyDlg message handlers

BOOL CMemSpyyDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	// Add "About..." menu item to system menu.

	// IDM_ABOUTBOX must be in the system command range.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != NULL)
	{
		CString strAboutMenu;
		strAboutMenu.LoadString(IDS_ABOUTBOX);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon


	FindPID();

	m_nTimer = SetTimer(1, 1000, 0);

	return TRUE;  // return TRUE  unless you set the focus to a control
}

void CMemSpyyDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialog::OnSysCommand(nID, lParam);
	}
}

// If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.

void CMemSpyyDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialog::OnPaint();
	}
}

// The system calls this function to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CMemSpyyDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}


void CMemSpyyDlg::SetLargestFree( unsigned long inLargestFree )
{
	char msg[255];
	sprintf( msg, "%4.2fMB", (double) ( (double) inLargestFree / 1024.0 / 1024.0) );
	m_LargestCtl.SetWindowText( msg );
}

void CMemSpyyDlg::SetTotalFree( unsigned long inTotalFree )
{
	char msg[255];
	sprintf( msg, "%4.2fMB", (double) ((double)inTotalFree / 1024.0 / 1024.0 ) );
	m_TotalFreeCtl.SetWindowText( msg );
}

void CMemSpyyDlg::SetTotalDLL( unsigned long inTotalDLL )
{
	char msg[255];
	sprintf( msg, "%4.2fMB", (double) ( (double)inTotalDLL / 1024.0 / 1024.0 ) );
	m_TotalDllCtl.SetWindowText( msg );
}

void CMemSpyyDlg::UpdateFromInjection()
{
	char tmpStr[255];

	sprintf( tmpStr, "%4.2fMB", g_UncommittedSize / 1024 / 1024 );

	m_ProcessHeapCtl.SetWindowText(tmpStr);

}

void CMemSpyyDlg::UpdateTotalVirtual()
{
	MEMORYSTATUSEX statex;
	statex.dwLength = sizeof (statex);
	GlobalMemoryStatusEx (&statex);

	char tmpStr[255];
	sprintf( tmpStr, "%4.2fMB", (double) ( (double) statex.ullTotalVirtual / 1024.0 / 1024.0) );

	m_TotalVirtSpaceCtl.SetWindowText( tmpStr );

}

void CMemSpyyDlg::OnTimer(UINT nIDEvent)
{
	m_OpenGLControl.ForceDraw();

	SetProcessHeapPtr( m_PID );
	UpdateTotalVirtual();
	UpdateFromInjection();

	CDialog::OnTimer(nIDEvent);
}

void CMemSpyyDlg::OnEnUpdatePid()
{
	// TODO:  If this is a RICHEDIT control, the control will not
	// send this notification unless you override the CDialog::OnInitDialog()
	// function to send the EM_SETEVENTMASK message to the control
	// with the ENM_UPDATE flag ORed into the lParam mask.

	UpdateData();

	m_OpenGLControl.SetPID( m_PID );
	m_PIDNameCtl.SetWindowText( GetEXENameByPID( m_PID ) );
}

void CMemSpyyDlg::FindPID()
{
	DWORD pids[4096];
	DWORD pidsSize = 0;
	DWORD numPids = 0;

	EnumProcesses(pids,sizeof(pids),&pidsSize);

	numPids = pidsSize / sizeof(DWORD);
	for (unsigned long x = 0; x < numPids; x++)
	{
		char pathName[255];

		HANDLE process = OpenProcess( PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,false, pids[x] );
		GetProcessImageFileName( process, pathName, 255);

		CString str = CString(pathName);

		CString strExeName = str.Right( str.GetLength() - str.ReverseFind( '\\') - 1 );

	}

}

CString CMemSpyyDlg::GetEXENameByPID( DWORD inPID )
{
	char pathName[255];

	HANDLE process = OpenProcess( PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,false, inPID );
	GetProcessImageFileName( process, pathName, 255);
	CString str = CString(pathName);
	
	CString strExeName = str.Right( str.GetLength() - str.ReverseFind( '\\') - 1 );

	return strExeName;
}

void CMemSpyyDlg::SetProcessHeapPtr( DWORD inPID)
{
	m_InjecteeProcess = OpenProcess( PROCESS_CREATE_THREAD | PROCESS_QUERY_INFORMATION | PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ,false, inPID );

	if ( m_InjecteeProcess )
	{
		g_Initialized = true;
		InjectDll( m_InjecteeProcess );
		if ( !m_Thread )
			m_InjecteeProcess = NULL;

		Cleanup();
	}

}

void CMemSpyyDlg::UpdateXY( double inX, double inY )
{
	char tmpStr[255];
	sprintf( tmpStr, "%4.2f", inX);
	m_XCoordCtl.SetWindowText(tmpStr);
	sprintf( tmpStr, "%4.2f", inY);
	m_YCoordCtl.SetWindowText(tmpStr);
}

void CMemSpyyDlg::UpdateStatus( PVOID inBaseAddress )
{
	// Get a handle to the process.
	HANDLE process = OpenProcess( PROCESS_QUERY_INFORMATION|PROCESS_VM_READ,false, m_PID );

	SIZE_T retVal = 0;
	MEMORY_BASIC_INFORMATION memBlock;
	retVal = VirtualQueryEx(process, inBaseAddress,&memBlock,sizeof(memBlock));
	if (retVal)
	{
		if (memBlock.State & MEM_FREE)
		{
			char tmpBuffer[255];
			sprintf( tmpBuffer, "Base: 0x%p Size: %4.2fMB   Description: Free", inBaseAddress, (double)memBlock.RegionSize/1024/1024);
			m_StatusCtl.SetWindowText( tmpBuffer );
		}
		else
		{
			char mappedFile[255];
			char tmpBuffer[255];
			if ( GetMappedFileName( process, inBaseAddress, mappedFile, 255) )
			{
				sprintf( tmpBuffer, "Base: 0x%p Size: %4.2fMB   Description: %s", inBaseAddress, (double) memBlock.RegionSize/1024/1024, strrchr(mappedFile, '\\') + 1 );
				m_StatusCtl.SetWindowText( tmpBuffer );
			}
			else
			{
				sprintf( tmpBuffer, "Base: 0x%p Size: %4.2fMB   Description: committed", inBaseAddress, (double) memBlock.RegionSize/1024/1024 );
				m_StatusCtl.SetWindowText( tmpBuffer );
			}
		}
	}
}

// I Tried splitting apart the injection and the 'de-jection', but I had trouble putting the loop inside the dll
// it seemed to block SOME of the injected apps.  Not sure why -- pj

//-----------------------------------------------
// InjectDll
// Notice: Loads "LibSpy.dll" into the remote process
//		   (via CreateRemoteThread & LoadLibrary)
//
void CMemSpyyDlg::InjectDll( HANDLE hProcess )
{
	HMODULE hKernel32 = ::GetModuleHandle("Kernel32");


	// Get full path of "LibSpy.dll"
	if( !GetModuleFileName( NULL,m_szLibPath,_MAX_PATH) )
		return;
	strcpy( strstr(m_szLibPath,"MemSpyy.exe"),"LibSpy.dll" );


	// 1. Allocate memory in the remote process for szLibPath
	// 2. Write szLibPath to the allocated memory
	m_pLibRemote = ::VirtualAllocEx( hProcess, NULL, sizeof(m_szLibPath), MEM_COMMIT, PAGE_READWRITE );
	if( m_pLibRemote == NULL )
		return;
	::WriteProcessMemory(hProcess, m_pLibRemote, (void*)m_szLibPath,sizeof(m_szLibPath),NULL);


	// Load "LibSpy.dll" into the remote process 
	// (via CreateRemoteThread & LoadLibrary)
	m_Thread = ::CreateRemoteThread( hProcess, NULL, 0,	
		(LPTHREAD_START_ROUTINE) ::GetProcAddress(hKernel32,"LoadLibraryA"), 
		m_pLibRemote, 0, NULL );
	if( m_Thread == NULL )
		FreeInjectDll( hProcess );
}

void CMemSpyyDlg::FreeInjectDll( HANDLE hProcess )
{
	HMODULE hKernel32 = ::GetModuleHandle("Kernel32");

	::VirtualFreeEx( hProcess, m_pLibRemote, sizeof(m_szLibPath), MEM_RELEASE );

	if( m_hLibModule == NULL )
		return;


	// Unload "LibSpy.dll" from the remote process 
	// (via CreateRemoteThread & FreeLibrary)
	m_Thread = ::CreateRemoteThread( hProcess,
		NULL, 0,
		(LPTHREAD_START_ROUTINE) ::GetProcAddress(hKernel32,"FreeLibrary"),
		(void*)m_hLibModule,
		0, NULL );
	if( m_Thread == NULL )	// failed to unload
		return;

	::WaitForSingleObject( m_Thread, INFINITE );
	::GetExitCodeThread( m_Thread, &m_hLibModule );
	::CloseHandle( m_Thread );
}

void CMemSpyyDlg::DejectDll( HANDLE hProcess )
{	
	g_Break = true;

	::WaitForSingleObject( m_Thread, INFINITE );

	// Get handle of loaded module
	::GetExitCodeThread( m_Thread, &m_hLibModule );
	::CloseHandle( m_Thread );

	FreeInjectDll( hProcess );
}

void CMemSpyyDlg::Cleanup()
{
	if (m_InjecteeProcess)
	{
		DejectDll( m_InjecteeProcess );
		m_InjecteeProcess = NULL;
	}
}